#include <arpa/inet.h>
#include <net/ethernet.h>
#include <netdb.h>
#include <netinet/ether.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <netinet/ip_icmp.h>
#include <pthread.h>
#include <sys/time.h>
#include <unistd.h>
#include <signal.h>
#include <syslog.h>
#include <fcntl.h>


#include <iostream>
#include <map>
#include <list>
#include <string>

extern "C" {
#include <pcap.h>
}

using namespace std;

#include "sqldata.h"
#include "ip.h"
#include "sll.h"
#include "dns.h"
#include "prio.h"
#include "monlist.h"

// TODO: put this in the RTable struucture

static struct bpf_program compiledSecondaryFilter;
static bool	secondaryFilterCompiled = false;

static struct bpf_program compiledNoSaveFilter;
static bool	noSaveFilterCompiled = false;

int	FlushDataToDb(RTable  *acom, int sinceInSeconds) {

	int seen = 0;

	PrioRecordList	*crl = new PrioRecordList();
	struct timeval now;
	
	gettimeofday(&now, NULL);
	if (sinceInSeconds == 0) {
		acom->MakeListOfFlushItems(crl);
	} else {
		struct timeval tv;
		memcpy(&tv, &now, sizeof (tv));
		tv.tv_sec -= sinceInSeconds;
		acom->MakeListOfFlushItems(crl, &tv);
	}
	for (PrioRecordListIter iter = crl->begin(); iter != crl->end(); ) {
		PrioRecordListIter prev = iter;
		CRecord *cr = static_cast<CRecord *>(*iter++);

		if ((cr->srcSerial_m = acom->FindHostFromIP(cr->ip_src_m, cr->llSrc_m)) == -1 ||
				(cr->dstSerial_m = acom->FindHostFromIP(cr->ip_dst_m, cr->llDst_m)) == -1) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "FindHostFromIP failed stopped flush list processing");
			break;
		}
		if (acom->SaveRecord(cr, acom->Get_InterfaceName(), sinceInSeconds == 0 ? &now : 0) == true) {
			seen++;
			crl->erase(prev);
			delete cr;
		} else {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "SaveRecord failed stopped flush list processing");
			break;
		}
	}
	delete crl;
	return seen;
}

static void	SetSignals(bool allowControlC)
{
	sigset_t  blockedSigs;

	/* Get the full signal set */
	sigfillset(&blockedSigs);

	if (allowControlC)
		sigdelset(&blockedSigs, SIGINT);
	/* Block all signals (except, possibly, 34) */
  pthread_sigmask(SIG_BLOCK, &blockedSigs, NULL);
}

static void    *SaveDataThread(void *arg)
{
	RTable	*rt = (RTable *)arg;

	SetSignals(false);
	for (;;) {
		for (int ii = 0; ii < rt->Get_SleepTime(); ii++) {
			sleep(1);
			pthread_testcancel();
		}
		FlushDataToDb(rt, rt->Get_Age());
		rt->Get_DnsCache()->FlushOld();
	}
}

static void    *FullFlushyThread(void *arg)
{
	RTable	*rt = (RTable *)arg;

	for (;;) {
		time_t	ti;
		time(&ti);

		if (ti >= rt->nextFlush_) {
			rt->nextFlush_ += rt->fullFlushInterval_;
			//cout << "Periodic flush" << ctime(&ti);
			int ii = FlushDataToDb(rt, 0);
			// syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "%d items flushed by periodic flusher", ii);
			rt->Get_DnsCache()->FlushOld();
		} 
		sleep(1);
		pthread_testcancel();

	}
}


u_char *handle_IP(u_char *args, const struct pcap_pkthdr *pkthdr, const u_char * packet, int hrdSize,
	const char *llSrc, const char *llDst, bool save)
{
	const u_char *payload = packet + hrdSize;
	struct ip	*iph = (struct ip *)payload;
	int unsigned payload_offset;
	int unsigned payload_length;
	int		port_src;
	int		port_dst;
		
	if (iph->ip_v != 4) {
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "unsupported IP version %d", iph->ip_v);
		return NULL;
	}
	switch (iph->ip_p) {
	case IPPROTO_ICMP: {
			struct icmphdr *icmp = (struct icmphdr *)(payload + (iph->ip_hl << 2));
			payload_offset = (sizeof (struct icmphdr) + (iph->ip_hl << 2) );
			payload_length = 0;
			port_src = 0;
			port_dst = 0;
		}
		break;
	case IPPROTO_TCP: {
			struct tcphdr *tcp = (struct tcphdr *)(payload + (iph->ip_hl << 2)); 
			payload_offset = ((iph->ip_hl << 2) + (tcp->doff << 2));
			payload_length = (unsigned int) ntohs(iph->ip_len) - ((iph->ip_hl << 2) + (tcp->doff << 2));
			port_src = tcp->source;
			port_dst = tcp->dest;
		}
		break;
	case IPPROTO_UDP: {
			struct udphdr *udp = (struct udphdr *)(payload + (iph->ip_hl << 2));
			payload_offset = (sizeof (struct udphdr) + (iph->ip_hl << 2) );
			payload_length = htons(udp->len) - sizeof (struct udphdr);
			port_src = udp->source;
			port_dst = udp->dest;
		}
		break;
	default:
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "Unsupported protocol %d", (int)iph->ip_p);
		return NULL;
	}
	CRecord	cr(iph, port_src, port_dst, pkthdr->len, llSrc, llDst,
				save ? string((const char *)payload + payload_offset, payload_length) : string(""));

	((RTable *)args)->Accumulate(&cr, pkthdr->ts);
	return NULL;
}


#ifndef ETHER_HDRLEN
#define ETHER_HDRLEN 14
#endif

void my_callback(u_char *args, const struct pcap_pkthdr *pkthdr, const u_char *packet)
{
	RTable *rt = (RTable *)args;
	u_int caplen = pkthdr->caplen;
	u_int length = pkthdr->len;
	bool save = false;

//	FlushDataToDb(rt, rt->Get_Age());
//	rt->Get_DnsCache()->FlushOld();

	if (secondaryFilterCompiled) {
		if (bpf_filter(compiledSecondaryFilter.bf_insns, (u_char *)packet, pkthdr->len, pkthdr->caplen) != 0) {
			save = true;
		}
	}
	if (noSaveFilterCompiled && 
				bpf_filter(compiledNoSaveFilter.bf_insns, (u_char *)packet, pkthdr->len, pkthdr->caplen) != 0) {
		return;
	}
	switch (rt->Get_DataLinkType()) {
	case DLT_PPP:
		break;

	case DLT_LINUX_SLL: {
			const struct sll_header *sllp = (const struct sll_header *)packet;
			u_short ether_type = ntohs(sllp->sll_protocol);
			switch (ether_type) {
			case ETHERTYPE_IP: {
					handle_IP(args, pkthdr, packet, SLL_HDR_LEN, NULL, NULL, save);
					break;
				}
			default:
				syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "unhandled ethertype: 0");
				break;
			}
		}
		break;


	case DLT_EN10MB: {
			const struct ether_header *eptr = (const struct ether_header *)packet;
			u_int16_t type = ntohs(eptr->ether_type);
			char	llSrc[80];
			char	llDst[80];
			if (caplen < ETHER_HDRLEN) {
				syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "Packet length less than ethernet header length");
				return;
			}
			switch (type) {
			case 0:
				syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "unhandled ethertype: 0");
				break;
			case ETHERTYPE_IP:
				ether_ntoa_r((struct ether_addr *)eptr->ether_shost, llSrc);
				ether_ntoa_r((struct ether_addr *)eptr->ether_dhost, llDst);
				handle_IP(args, pkthdr, packet, sizeof (struct ether_header), llSrc, llDst, save);
				break;
			case ETHERTYPE_ARP:
				break;
			case ETHERTYPE_REVARP:
				break;
			default:
				break;
			}
		}
		break;
	default:
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "unknown data link layer type %d",  rt->Get_DataLinkType());
		break;
	}
}

RTable   *rt;
pthread_t	writeThread;
pthread_t	flushyThread;
 
static void	FinalFlush()
{
	int retval;
	syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "got signal waiting for DNS records to be processed and fushed");
	if (flushyThread) {
		pthread_cancel(flushyThread);
		retval = pthread_join(flushyThread, NULL);
		if ( retval < 0) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "pthread_join(flushyThread, .. failed");
			exit(EXIT_FAILURE);
		}
	}
	pthread_cancel(writeThread);
	
	retval = pthread_join(writeThread, NULL);
	if ( retval < 0) {
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "pthread_join(writeThread, .. failed");
		exit(EXIT_FAILURE);
	}
	if (FlushDataToDb(rt, 0)) {
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "committed further records");
	} else {
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "not committed further records");
	}
}

static void termination_handler(int signum)
{
	FinalFlush();

	struct pcap_stat ps;
	if (rt->Get_descr() != 0) {
		if (pcap_stats(rt->Get_descr(), &ps) < 0) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "pcap_stats failed");
			exit(1);
		}
	}
	FinalFlush();
	syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "%d received and %d dropped", ps.ps_recv, ps.ps_drop);
	exit(0);
}

bool	CompileIfNeeded(char *filterText, bool filterCompiled,
			pcap_t *descr,  bpf_u_int32 netp, struct bpf_program& filter,
			char *typeAsText)
{
	if (filterText != 0 && !filterCompiled) {
		if (pcap_compile(descr, &filter, filterText, 0, netp) == -1) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "Error calling pcap_compile for %s", typeAsText);
			exit(1);
		} else {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "%s compiled OK", typeAsText);
		}
		return true;
	} else {
		return false;
	}
}

int main(int argc, char *argv[])
{
    int nerr;

	char *dev = (char *)0;
	char errbuf[PCAP_ERRBUF_SIZE];

	int	opt;
	int	sltime = 10;
	int	numPackets = 10;
	int	age = 120;
	char	*connectionString = "usage";
	char	*filter = 0;
	char	*secondaryFilter = 0;
	int		dnsExpiryTime = 120;
	int	backGround = false;
	int	ifDownSleep = 10;
	int	fullFlushOnMinDiv = 0;
	char	*noSaveFilter = 0;

	while ((opt = getopt(argc, argv, "F:S:bx:s:p:c:f:e:d:i:a:n:")) != -1) {
		switch (opt) {
		case 'F':
			fullFlushOnMinDiv = atoi(optarg);
			break;
		case 'S':
			ifDownSleep = atoi(optarg);
			break;
		case 'b':
			backGround = true;
			break;
		case 'x':
			dnsExpiryTime = atoi(optarg);
			break;
		case 's':
			sltime = atoi(optarg);
			break;
		case 'p':
			numPackets = atoi(optarg);
			break;
		case 'c':
			connectionString = optarg;
			break;
		case 'f':
			filter = optarg;
			break;
		case 'e':
			secondaryFilter = optarg;
			break;
		case 'd':
		case 'i':
			dev = optarg;
			break;
		case 'a':
			age = atoi(optarg);
			break;
		case 'n':
			noSaveFilter = optarg;
			break;
		default:
			cerr << "invalid usage" << endl;
			cerr << "-F seconds - fully flush the usage list with this seconds divisor" << endl;
			cerr << " for example, 10 - every 10 minutes but at 1:10, 1:20 etc" << endl;
			cerr << "-b - fork in background and detatch from console" << endl;
			cerr << "-x dns expiry time: the time an entry lives int he dns cache" << endl;
			cerr << "-s sleep_time: the time the write sleeps for" << endl;
			cerr << "-p num_packets: the number of packets to collect, 0 for forever" << endl;
			cerr << "-f filter: pcap filter argument" << endl;
			cerr << "-e seconday filter: filter records to save the whole data" << endl;
			cerr << "-d or -i: device (or interface) to capure on" << endl;
			cerr << "-a: age of to/from packets inactivity before writing to db" << endl;
			cerr << "-c: connection string" << endl;
			cerr << "-S: number of seconds to sleep if interface is down" << endl;
			cerr << "-n: filter specification for flows not to record" << endl;
			exit(1);
		}
	}
	if (dev == (char *)0) {
		if ((dev = pcap_lookupdev(errbuf)) == (char *)0) {
			cerr << "pcap_lookupdev failed: " << errbuf << endl;
			exit(1);
		}
	}
	if (backGround) {
		int fval =fork();
		if (fval < 0) {
			perror("fork");
			exit(1); /* fork error */
		} else if (fval > 0) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "forked child %d", fval);
			exit(0); /* parent exits */
		}
		setsid();
	}

	syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "capturing on interface '%s'", dev);

	bpf_u_int32 maskp;    /* subnet mask   */
	bpf_u_int32 netp;     /* ip   */           
	pcap_t *descr;			/* pcap descriptor */

	if (pcap_lookupnet(dev, &netp, &maskp, errbuf) < 0) {
		cerr << "pcap_lookupnet(): " << errbuf << endl;
		exit(1);
	}

	DnsCache	dnsCache(dnsExpiryTime);
	rt = new RTable(sltime, &dnsCache, age, fullFlushOnMinDiv, dev, connectionString);

#if 1
   if (pthread_create(&writeThread, NULL, SaveDataThread, (void *)rt) < 0) {
		perror("pthread_create");
		exit(1);
    }
#endif
	if (fullFlushOnMinDiv != 0) {
		if (pthread_create(&flushyThread, NULL, FullFlushyThread, (void *)rt) < 0) {
			perror("pthread_create for flushy");
			exit(1);
		}
	}
	signal(SIGINT, termination_handler);
	signal(SIGTERM, termination_handler);

	bool	linkFilterCompiled = false;

	for (;;) {
		struct bpf_program compiledFilter;
		
		if ((descr = pcap_open_live(dev, BUFSIZ, 1, -1, errbuf)) == (pcap_t *)0) {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "pcap_open_live(): %s", errbuf);
			if (numPackets != 0)
				exit(1);
		} else {
			syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "capturing datalink type %d", pcap_datalink(descr));
			rt->Set_tcapInfo(descr);

			if (filter != 0 && !linkFilterCompiled) {
				if (pcap_compile(descr, &compiledFilter, filter, 0, netp) == -1) {
					syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "Error calling pcap_compile");
					exit(1);
				}
				linkFilterCompiled = true;
			}
			if (secondaryFilter != 0 && !secondaryFilterCompiled) {
				if (pcap_compile(descr, &compiledSecondaryFilter, secondaryFilter, 0, netp) == -1) {
					syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "Error calling pcap_compile for secondary");
					exit(1);
				} else {
					syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "secondary filter compiled OK");
					
				}
				secondaryFilterCompiled = true;
				
			}

			noSaveFilterCompiled = CompileIfNeeded(noSaveFilter, noSaveFilterCompiled, 
					descr, netp, compiledNoSaveFilter,
					"No save filter expression");

			if (filter != 0) {
				/* set the compiled program as the filter */
				if (pcap_setfilter(descr, &compiledFilter) == -1) {
					cerr << "Error setting filter" << endl;
					exit(1); 
				} 
			}

			pcap_loop(descr, numPackets, my_callback, (u_char *)rt);
			if (numPackets != 0)
				break;
		}
		syslog(LOG_MAKEPRI(LOG_DAEMON, LOG_INFO), "down .. sleeping %d", ifDownSleep);
		sleep(ifDownSleep);
	}
	FinalFlush();

	return 0;
}


